#include "StdAfx.h"

#include "Encode.h"
BOOL	DumpFlatBinData(BYTE* stm32, DWORD from, DWORD to, LPCTSTR name)
{
	_TCHAR	flatFile[MAX_PATH] = _T("");
	size_t len = _tcslen(name);
	_tcscpy(flatFile,name);
	flatFile[len-4] = _T('.');
	flatFile[len-3] = _T('r');
	flatFile[len-2] = _T('o');
	flatFile[len-1] = _T('m');
	FILE* fp = _tfopen(flatFile,_T("w+"));
	if(fp == NULL){ return FALSE; }
	size_t fileLen = to - from;
	size_t actLen = fwrite(stm32+(from-FLASH_BASE),1,fileLen,fp);
	fclose(fp);
	if(actLen != fileLen){
		printf("Fail generate binary file!\n");
		return FALSE;
	}
	return TRUE;
}


BOOL	DumpEncodeBinData(STM32_BIN& binData,LPCTSTR outputFile,size_t &blockSize)
{
	FILE* fp = _tfopen(outputFile,_T("wb+"));
	if(fp == NULL){ return FALSE; }
	/// Do not contain unused blocks
	size_t fileLen = sizeof(STM32_BIN) 
		- sizeof(AppBlock)*(AppMaxBlock - blockSize);

	BYTE* p = (BYTE*)&binData;
	binData.fileLen = fileLen;
	binData.crcFile = 
	CalcCrc32Rev( ((BYTE*)&binData) + 4,  
	// Not include the CRC valu
		fileLen - 4);

	size_t actLen = fwrite(&binData,1,fileLen,fp);
	fclose(fp);
	if(actLen != fileLen){
		printf("Fail generate binary file!\n");
		return FALSE;
	}
	return TRUE;
}

BOOL	LoadHexForSTM32(LPCTSTR filename, STM32_BIN	&STM32BinData, BYTE* STM32_Flash)
{
	try{
		CHexFile	hex(filename);
	if(!hex.LineCount){
		printf("Fail to load hex file!\n");
		return FALSE;
	}
	memset(STM32BinData.resrved,0,sizeof(STM32BinData.resrved));
	memcpy(STM32BinData.resrved,"Thanks for use PQLabs' Product!",sizeof(STM32BinData.resrved));
	memset(STM32_Flash,0xff,FLASH_SIZE);
	DWORD	linearAddress = 0;
	DWORD	lastFlashAddr = 0;
	DWORD	startFlashAddr = AppParamAddr;
	for(int i =0; i<hex.LineCount;i++){
		CHexFileLine& line = hex.Lines[i];
		switch(line.RecordType){
			case t_Linear:
				linearAddress = 
			(((DWORD)line.Records[0])<<24) + 
			(((DWORD)line.Records[1])<<16);
				break;
			case t_Data:
				if(line.RecordAddress + linearAddress >= FLASH_BASE){
					DWORD addr = line.RecordAddress + linearAddress - FLASH_BASE;
					if(addr<FLASH_SIZE){
                        //for(int i=0;i<line.RecordNumber;i++){
                        //    if(addr+i+FLASH_BASE>=0x08003E00 && addr+i+FLASH_BASE<AppBaseAddr){
                        //    }else{
                        //        STM32_Flash[addr+i] = line.Records[i];
                        //    }
                        //}
						memcpy(&STM32_Flash[addr],
						line.Records,
						line.RecordNumber );
						if(lastFlashAddr < addr+line.RecordNumber-1+FLASH_BASE){
							lastFlashAddr = addr+line.RecordNumber-1+FLASH_BASE;
						}
					}
				}
				break;
			default:
				break;
		}
	}
	printf("Range: 0x%08X - 0x%08X\nSize: %d bytes\n",
		startFlashAddr,lastFlashAddr,lastFlashAddr-startFlashAddr);
	STM32BinData.header.appStartAddr = startFlashAddr &0xFFFFFFFC;
	STM32BinData.header.appSize = (lastFlashAddr + 1 - startFlashAddr + 3)/4;
	STM32BinData.header.appCrc32 = CrcGen_STM32(
		(unsigned long*)&STM32_Flash[STM32BinData.header.appStartAddr - FLASH_BASE],
		STM32BinData.header.appSize);
	return TRUE;
	}
	catch(...){
		return FALSE;
	}
}

BOOL IsBlockEmpty(DWORD* flash)
{
	for(size_t i=0;i<BLOCK_SIZE;i++){
		if(flash[i]!=0xFFFFFFFF){
			return FALSE;
		}
	}
	return TRUE;
}


BOOL CreateBlockData(STM32_BIN	&STM32BinData, BYTE* STM32_Flash, size_t &blockSize)
{
	MakeCrc32TableRev();
	DWORD crcSN = CalcCrc32Rev(
		(const BYTE*)STM32BinData.wstrSerialNumber,
		24);
	srand(crcSN);
	for(int i=0;i<16;i++){
		STM32BinData.cipher[i] = rand();
#ifdef	DEBUG
		if(i==0){printf("\n");}
		printf("%02x, ",STM32BinData.cipher[i]);
		if(i==7||i==15){printf("\n");}
#endif
	}
	for(int i=0;i<16;i++){
		STM32BinData.random[i] = rand();
#ifdef	DEBUG
		if(i==0){printf("\n");}
		printf("%02x, ",STM32BinData.random[i]);
		if(i==7||i==15){printf("\n");}
#endif
	}

	size_t i=0;
	size_t curBlock = 0;
	DWORD* appStartAddr = (DWORD*)&STM32_Flash[STM32BinData.header.appStartAddr - FLASH_BASE];
	while(i<STM32BinData.header.appSize){
		if(curBlock>=AppMaxBlock){
			printf("Block out of range!\n");
			return FALSE;
		}
		if(IsBlockEmpty(appStartAddr+i)){
			i+=BLOCK_SIZE;
			continue;
		}
		AppBlock& block = STM32BinData.AppBlock[curBlock];
		block.address = (appStartAddr-(DWORD*)STM32_Flash + i)*4 + FLASH_BASE;
		block.size = BLOCK_SIZE;
		memcpy(block.data,appStartAddr+i,BLOCK_SIZE*4);
		curBlock++;
		i+=BLOCK_SIZE;
	}
///////////////////////////////////////////////////////////////////////////////////////
//// Add some useless blocks, this will make the encoded data with different length
	size_t mixLen = (size_t)rand();
	mixLen %= 30; /// allow 30 addition blocks
	while(mixLen){
		AppBlock& block = STM32BinData.AppBlock[curBlock];
		block.address = (unsigned long)rand();
		while(block.address>=MinProgAddr){
			block.address-= MinProgAddr;
		}
		block.size = BLOCK_SIZE;
		for(size_t i=0;i<block.size;i++){
			block.data[i] = (unsigned long)rand();
		}
		curBlock++;
		mixLen--;
	}
/////////////////////////////////////////////////////////////////////////////////////////
	STM32BinData.header.appBlockSize = curBlock;
	blockSize = curBlock;
	return TRUE;
}

void	SwapBlock(STM32_BIN	&STM32BinData, size_t &blockSize)
{
	DWORD crcSN = CalcCrc32Rev(
		(const BYTE*)STM32BinData.wstrSerialNumber,
		24);
	srand(crcSN);
	for(size_t i =0 ; i< 1000 ; i++){
		size_t p1 = rand()%blockSize;
		size_t p2 = rand()%blockSize;
		if(p1!=p2){
			AppBlock temp;
			memcpy(&temp,&STM32BinData.AppBlock[p1],sizeof(AppBlock));
			memcpy(&STM32BinData.AppBlock[p1],&STM32BinData.AppBlock[p2],sizeof(AppBlock));
			memcpy(&STM32BinData.AppBlock[p2],&temp,sizeof(AppBlock));
		}
	}
}

BOOL	EncodeBlockData(STM32_BIN	&STM32BinData, size_t &blockSize)
{
	SwapBlock(STM32BinData,blockSize);
	BYTE chain[16];
	memcpy(chain,STM32BinData.cipher,16);
	aesEncInit();
	aesEncrypt((BYTE*)&STM32BinData.header,chain);
	for(size_t i=0;i<blockSize;i++){
		BYTE* block = (BYTE*)&STM32BinData.AppBlock[i];
		for(size_t j=0;j<sizeof(AppBlock);j+=16){
			aesEncrypt(block+j,chain);
		}
	}
	return TRUE;
}

//////////////////////////////////////////////////////////////////
/// Old main codes
#include "..\KeyMan\KeyDataBase.h"
void Usage()
{
	printf(	"Usage:\n"
		"GenEncCode <HexfileName> [Serial Number]\n"
		//"GenEncCode <HexfileName>\n"
		);
}
extern unsigned char kTable[32];
void	MakeKey(unsigned char table[32], ONE_KEY_DATA &key, BYTE* Keys)
{
	DWORD	keyType = key.keyType[0] == '1' ? 1 : 0;
	WORD keySign;
	BYTE	cid[12];
	{
		char	tmp[256] = "";
		memcpy(tmp,key.strKey,4);
		DWORD tp;
		::sscanf(tmp,"%04X",&tp);
		keySign = (WORD)tp;
	}
	{
		char	tmp[256] = "";
		for(size_t i=0;i<12;i++){
			memcpy(tmp,key.strChipID+i*2,2);
			DWORD tp;
			sscanf(tmp,"%02X",&tp);
			cid[i] = (BYTE)tp;
		}
	}
	for(size_t i=0;i<12;i++){
		table[i] =  (keyType == KT_XOR) ? 
		(Keys+keySign*32)[i]^cid[i] : (Keys+keySign*32)[i];
	}
	for(size_t i=12;i<32;i++){
		table[i] =  (Keys+keySign*32)[i];
	}
}

//////////////////////////////////////////////////////////////////
/// argc    must be 4
/// argv[0] app path
/// argv[1] hex file path
/// argv[2] serial number
/// argv[3] output file name
int CreateEncodeBinFile(int argc, _TCHAR* argv[],CKeyDataBase& keyDB, BYTE* keys)
{
	ONE_KEY_DATA keyData;
	if(!keyDB.GetKey(argv[2],&keyData)){
		return -1;
	}
	MakeKey(kTable,keyData, keys);
	
	_TCHAR	outputFile[MAX_PATH] = _T("");
	size_t blockSize = 0;
	STM32_BIN	STM32BinData;
	BYTE	STM32_Flash[FLASH_SIZE];
	if(argc!=4){
		Usage();
		return -1;
	}
	if(argc<3){
		Usage();
		return -1;
	}
	if(argc>3){
		_tcscpy(outputFile,argv[3]);
	}else{
		size_t len = _tcslen(argv[1]);
		_tcscpy(outputFile,argv[1]);
		outputFile[len-4] = _T('.');
		outputFile[len-3] = _T('b');
		outputFile[len-2] = _T('i');
		outputFile[len-1] = _T('n');
	}
	if(!LoadHexForSTM32(argv[1],STM32BinData,STM32_Flash)){
		printf("Fail to load hex file for STM32!\n");
		return -1;
	}
	//DumpFlatBinData(STM32_Flash,0x08003C00,0x08003C00+STM32BinData.header.appSize*4-1,argv[1]);

	PDEVICE_SERIAL pDevParam = (PDEVICE_SERIAL)( (BYTE*)STM32_Flash + ((DWORD)DevSerial() - FLASH_BASE));
	PDEVICE_VERSION pFirmParam = (PDEVICE_VERSION)( (BYTE*)STM32_Flash + ((DWORD)DevVersion() - FLASH_BASE));
	if(pDevParam->valid == VALID_SIGN){
		STM32BinData.bcdHardwareVersion = pDevParam->hardVersion;
		for(size_t i=0;i<12;i++){
			//// Compare the serial number
			if(pDevParam->wstrSerial[i] != _T('*')){
				if(pDevParam->wstrSerial[i] != argv[2][i]){
					wprintf(L"Serial Number Missmatch! Ignore!\nRequire: %s,",pDevParam->wstrSerial);
					_tprintf(_T("Actual: %s\n"),argv[2]);
					return -1;
				}
			}
		}
		unsigned char toStr[32]="0123456789ABCEDF";
		unsigned char hw[4] = {(BYTE)((pDevParam->hardVersion&0xF000)>>12),
					  (BYTE)((pDevParam->hardVersion&0x0F00)>>8),
					  (BYTE)((pDevParam->hardVersion&0x00F0)>>4),
					  (BYTE)((pDevParam->hardVersion&0x000F)>>0)};
		ONE_KEY_DATA keyData;
		keyDB.GetKey(argv[2],&keyData);
		for(size_t i=0;i<4;i++){
			hw[i] = toStr[hw[i]];
			if(hw[i]!='A'){
				if(hw[i] != keyData.strHW[i]){
					char tp[16] = {keyData.strHW[0],keyData.strHW[1],'.',keyData.strHW[2],keyData.strHW[3]};
					printf("Hardware version missmatch! Ignore! \nRequire: %02X.%02X  Actual: %s\n",
						(pDevParam->hardVersion>>8)&0xFF,(pDevParam->hardVersion)&0xFF,
						tp);
					return -1;
				}
			}
		}
		
		wprintf(L"Serial Number:    %s, Hardware version: %X.%02X\n",pDevParam->wstrSerial
			,(pDevParam->hardVersion>>8)&0xFF,(pDevParam->hardVersion)&0xFF);
	}else{
		printf("No hardware specified!\n");
		return -1;
	}
	if(pFirmParam->valid == VALID_SIGN){
		memset(STM32BinData.wstrFirmwareVersion,0,sizeof(STM32BinData.wstrFirmwareVersion));
		for(size_t i=0;i<10;i++){
		STM32BinData.wstrFirmwareVersion[i] = pFirmParam->wstrVersion[i];
	}
		wprintf(L"Firmware version: %s\n",pFirmParam->wstrVersion);
	}else{
		printf("No firmware version infomation!\n");
		return -1;
	}
	memset(STM32BinData.wstrSerialNumber,0,sizeof(STM32BinData.wstrSerialNumber));
	for(size_t i=0;i<12;i++){
		STM32BinData.wstrSerialNumber[i] = (argv[2])[i];
	}
	CreateBlockData(STM32BinData,STM32_Flash,blockSize);
	EncodeBlockData(STM32BinData, blockSize);
	_tprintf(_T("The output file is place at:\n%s\n"),outputFile);
	DumpEncodeBinData(STM32BinData,outputFile,blockSize);
	printf("Start verify the output file...\n");
	if(!DecodeBin(outputFile)){
		printf("Data verify fail!\n");
		return -1;
	}else{
		printf("Output file verify OK!\n");
	}
	return 0;
}

